{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ab8e4db4",
   "metadata": {},
   "source": [
    "# conv2d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "feacf63a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7c243914",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Conv2DWorkload(batch=1, height=56, width=56, in_filter=64, out_filter=64, hkernel=3, wkernel=3, hpad=1, wpad=1, hstride=1, wstride=1)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-27 21:00:43.792 INFO load_module /tmp/tmph0e_m_py/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU CONV2D TEST PASSED: Time cost = 0.00506093 sec/op, 45.6854 GOPS\n",
      "Conv2DWorkload(batch=1, height=56, width=56, in_filter=64, out_filter=128, hkernel=3, wkernel=3, hpad=1, wpad=1, hstride=2, wstride=2)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-27 21:00:46.097 INFO load_module /tmp/tmph0e_m_py/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU CONV2D TEST PASSED: Time cost = 0.00413532 sec/op, 27.9557 GOPS\n",
      "Conv2DWorkload(batch=1, height=56, width=56, in_filter=64, out_filter=128, hkernel=1, wkernel=1, hpad=0, wpad=0, hstride=2, wstride=2)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-27 21:00:47.487 INFO load_module /tmp/tmph0e_m_py/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU CONV2D TEST PASSED: Time cost = 0.000443068 sec/op, 28.9911 GOPS\n",
      "Conv2DWorkload(batch=1, height=28, width=28, in_filter=128, out_filter=128, hkernel=3, wkernel=3, hpad=1, wpad=1, hstride=1, wstride=1)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-27 21:00:49.527 INFO load_module /tmp/tmph0e_m_py/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU CONV2D TEST PASSED: Time cost = 0.00652784 sec/op, 35.4192 GOPS\n",
      "Conv2DWorkload(batch=1, height=28, width=28, in_filter=128, out_filter=256, hkernel=3, wkernel=3, hpad=1, wpad=1, hstride=2, wstride=2)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-27 21:00:52.580 INFO load_module /tmp/tmph0e_m_py/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU CONV2D TEST PASSED: Time cost = 0.00161277 sec/op, 71.6812 GOPS\n",
      "Conv2DWorkload(batch=1, height=28, width=28, in_filter=128, out_filter=256, hkernel=1, wkernel=1, hpad=0, wpad=0, hstride=2, wstride=2)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-27 21:00:54.705 INFO load_module /tmp/tmph0e_m_py/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU CONV2D TEST PASSED: Time cost = 0.000355448 sec/op, 36.1377 GOPS\n",
      "Conv2DWorkload(batch=1, height=14, width=14, in_filter=256, out_filter=256, hkernel=3, wkernel=3, hpad=1, wpad=1, hstride=1, wstride=1)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-27 21:00:57.986 INFO load_module /tmp/tmph0e_m_py/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU CONV2D TEST PASSED: Time cost = 0.00332815 sec/op, 69.4713 GOPS\n",
      "Conv2DWorkload(batch=1, height=14, width=14, in_filter=256, out_filter=512, hkernel=3, wkernel=3, hpad=1, wpad=1, hstride=2, wstride=2)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-27 21:01:05.167 INFO load_module /tmp/tmph0e_m_py/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU CONV2D TEST PASSED: Time cost = 0.00237562 sec/op, 48.6633 GOPS\n",
      "Conv2DWorkload(batch=1, height=14, width=14, in_filter=256, out_filter=512, hkernel=1, wkernel=1, hpad=0, wpad=0, hstride=2, wstride=2)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-27 21:01:09.774 INFO load_module /tmp/tmph0e_m_py/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU CONV2D TEST PASSED: Time cost = 0.000318768 sec/op, 40.296 GOPS\n",
      "Conv2DWorkload(batch=1, height=7, width=7, in_filter=512, out_filter=512, hkernel=3, wkernel=3, hpad=1, wpad=1, hstride=1, wstride=1)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-27 21:01:18.716 INFO load_module /tmp/tmph0e_m_py/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU CONV2D TEST PASSED: Time cost = 0.00399377 sec/op, 57.893 GOPS\n",
      "Conv2DWorkload(batch=1, height=56, width=56, in_filter=64, out_filter=64, hkernel=3, wkernel=3, hpad=1, wpad=1, hstride=1, wstride=1)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[21:01:19] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
      "2025-07-27 21:01:19.957 INFO load_module /tmp/tmp7sixkb3a/conv2d.o\n",
      "[21:01:19] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/runtime/profiling.cc:101: Warning: No timer implementation for ext_dev, using default timer instead. It may be inaccurate or have extra overhead.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VTA CONV2D TEST PASSED: Time cost = 0.0333152 sec/op, 6.9401 GOPS\n",
      "Conv2DWorkload(batch=1, height=56, width=56, in_filter=64, out_filter=128, hkernel=3, wkernel=3, hpad=1, wpad=1, hstride=2, wstride=2)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[21:01:21] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
      "2025-07-27 21:01:21.789 INFO load_module /tmp/tmp7sixkb3a/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VTA CONV2D TEST PASSED: Time cost = 0.0174461 sec/op, 6.62643 GOPS\n",
      "Conv2DWorkload(batch=1, height=56, width=56, in_filter=64, out_filter=128, hkernel=1, wkernel=1, hpad=0, wpad=0, hstride=2, wstride=2)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[21:01:22] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
      "2025-07-27 21:01:22.844 INFO load_module /tmp/tmp7sixkb3a/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VTA CONV2D TEST PASSED: Time cost = 0.00232785 sec/op, 5.518 GOPS\n",
      "Conv2DWorkload(batch=1, height=28, width=28, in_filter=128, out_filter=128, hkernel=3, wkernel=3, hpad=1, wpad=1, hstride=1, wstride=1)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[21:01:23] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
      "2025-07-27 21:01:24.249 INFO load_module /tmp/tmp7sixkb3a/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VTA CONV2D TEST PASSED: Time cost = 0.0328963 sec/op, 7.02848 GOPS\n",
      "Conv2DWorkload(batch=1, height=28, width=28, in_filter=128, out_filter=256, hkernel=3, wkernel=3, hpad=1, wpad=1, hstride=2, wstride=2)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[21:01:26] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
      "2025-07-27 21:01:27.038 INFO load_module /tmp/tmp7sixkb3a/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VTA CONV2D TEST PASSED: Time cost = 0.0164558 sec/op, 7.02523 GOPS\n",
      "Conv2DWorkload(batch=1, height=28, width=28, in_filter=128, out_filter=256, hkernel=1, wkernel=1, hpad=0, wpad=0, hstride=2, wstride=2)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[21:01:28] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
      "2025-07-27 21:01:28.841 INFO load_module /tmp/tmp7sixkb3a/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VTA CONV2D TEST PASSED: Time cost = 0.00212144 sec/op, 6.05488 GOPS\n",
      "Conv2DWorkload(batch=1, height=14, width=14, in_filter=256, out_filter=256, hkernel=3, wkernel=3, hpad=1, wpad=1, hstride=1, wstride=1)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[21:01:31] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
      "2025-07-27 21:01:31.732 INFO load_module /tmp/tmp7sixkb3a/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VTA CONV2D TEST PASSED: Time cost = 0.0324786 sec/op, 7.11888 GOPS\n",
      "Conv2DWorkload(batch=1, height=14, width=14, in_filter=256, out_filter=512, hkernel=3, wkernel=3, hpad=1, wpad=1, hstride=2, wstride=2)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[21:01:37] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
      "2025-07-27 21:01:37.950 INFO load_module /tmp/tmp7sixkb3a/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VTA CONV2D TEST PASSED: Time cost = 0.0164752 sec/op, 7.01693 GOPS\n",
      "Conv2DWorkload(batch=1, height=14, width=14, in_filter=256, out_filter=512, hkernel=1, wkernel=1, hpad=0, wpad=0, hstride=2, wstride=2)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[21:01:42] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
      "2025-07-27 21:01:42.659 INFO load_module /tmp/tmp7sixkb3a/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VTA CONV2D TEST PASSED: Time cost = 0.00197737 sec/op, 6.49603 GOPS\n",
      "Conv2DWorkload(batch=1, height=7, width=7, in_filter=512, out_filter=512, hkernel=3, wkernel=3, hpad=1, wpad=1, hstride=1, wstride=1)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[21:01:51] /media/pc/data/board/arria10/lxw/tasks/tvm-new/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement  required_alignment=256, provided_alignment=64\n",
      "2025-07-27 21:01:51.528 INFO load_module /tmp/tmp7sixkb3a/conv2d.o\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VTA CONV2D TEST PASSED: Time cost = 0.0322971 sec/op, 7.15889 GOPS\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import os\n",
    "\n",
    "import pytest\n",
    "import numpy as np\n",
    "from collections import namedtuple\n",
    "\n",
    "import tvm\n",
    "from tvm import te\n",
    "from tvm import relay\n",
    "from tvm import autotvm\n",
    "from tvm.contrib import utils\n",
    "from tvm.contrib.pickle_memoize import memoize\n",
    "from tvm import topi\n",
    "import tvm.topi.testing\n",
    "import vta\n",
    "from vta import program_fpga, reconfig_runtime\n",
    "import vta.testing\n",
    "from vta.testing import simulator\n",
    "\n",
    "\n",
    "Workload = namedtuple(\n",
    "    \"Conv2DWorkload\",\n",
    "    [\n",
    "        \"batch\",\n",
    "        \"height\",\n",
    "        \"width\",\n",
    "        \"in_filter\",\n",
    "        \"out_filter\",\n",
    "        \"hkernel\",\n",
    "        \"wkernel\",\n",
    "        \"hpad\",\n",
    "        \"wpad\",\n",
    "        \"hstride\",\n",
    "        \"wstride\",\n",
    "    ],\n",
    ")\n",
    "\n",
    "# Get batch info from env\n",
    "env = vta.get_env()\n",
    "\n",
    "# ResNet18 workloads\n",
    "resnet_wkls = [\n",
    "    # Workloads of resnet18 on imagenet\n",
    "    # ('resnet-18.C1',  Workload(env.BATCH, 224, 224, 3,   64,  7, 7, 3, 3, 2, 2)),\n",
    "    (\"resnet-18.C2\", Workload(env.BATCH, 56, 56, 64, 64, 3, 3, 1, 1, 1, 1)),\n",
    "    (\"resnet-18.C3\", Workload(env.BATCH, 56, 56, 64, 128, 3, 3, 1, 1, 2, 2)),\n",
    "    (\"resnet-18.C4\", Workload(env.BATCH, 56, 56, 64, 128, 1, 1, 0, 0, 2, 2)),\n",
    "    (\"resnet-18.C5\", Workload(env.BATCH, 28, 28, 128, 128, 3, 3, 1, 1, 1, 1)),\n",
    "    (\"resnet-18.C6\", Workload(env.BATCH, 28, 28, 128, 256, 3, 3, 1, 1, 2, 2)),\n",
    "    (\"resnet-18.C7\", Workload(env.BATCH, 28, 28, 128, 256, 1, 1, 0, 0, 2, 2)),\n",
    "    (\"resnet-18.C8\", Workload(env.BATCH, 14, 14, 256, 256, 3, 3, 1, 1, 1, 1)),\n",
    "    (\"resnet-18.C9\", Workload(env.BATCH, 14, 14, 256, 512, 3, 3, 1, 1, 2, 2)),\n",
    "    (\"resnet-18.C10\", Workload(env.BATCH, 14, 14, 256, 512, 1, 1, 0, 0, 2, 2)),\n",
    "    (\"resnet-18.C11\", Workload(env.BATCH, 7, 7, 512, 512, 3, 3, 1, 1, 1, 1)),\n",
    "]\n",
    "\n",
    "# FIXME: we need a custom clip operator to circumvent a pattern detection limitation\n",
    "@tvm.te.tag_scope(tag=topi.tag.ELEMWISE)\n",
    "def my_clip(x, a_min, a_max):\n",
    "    \"\"\"Unlike topi's current clip, put min and max into two stages.\"\"\"\n",
    "    const_min = tvm.tir.const(a_min, x.dtype)\n",
    "    const_max = tvm.tir.const(a_max, x.dtype)\n",
    "    x = te.compute(x.shape, lambda *i: tvm.te.min(x(*i), const_max), name=\"clipA\")\n",
    "    x = te.compute(x.shape, lambda *i: tvm.te.max(x(*i), const_min), name=\"clipB\")\n",
    "    return x\n",
    "\n",
    "\n",
    "def run_conv2d(env, remote, wl, target, check_correctness=True, print_ir=False, samples=4):\n",
    "\n",
    "    # Workload assertions\n",
    "    assert wl.hpad == wl.wpad\n",
    "\n",
    "    # Perform packing only if we are targeting the accelerator\n",
    "    if \"arm_cpu\" in target.keys:\n",
    "        data_pack = False\n",
    "        layout = \"NCHW\"\n",
    "        conv2d_fcompute = topi.arm_cpu.conv2d_nchw_spatial_pack\n",
    "        conv2d_fschedule = topi.arm_cpu.schedule_conv2d_nchw_spatial_pack\n",
    "    elif \"vta\" in target.keys:\n",
    "        data_pack = True\n",
    "        layout = \"NCHW%dn%dc\" % (env.BATCH, env.BLOCK_IN)\n",
    "        conv2d_fcompute = vta.top.conv2d_packed\n",
    "        conv2d_fschedule = vta.top.schedule_conv2d_packed\n",
    "\n",
    "    # Derive shapes depending upon packing\n",
    "    a_shape = (wl.batch, wl.in_filter, wl.height, wl.width)\n",
    "    w_shape = (wl.out_filter, wl.in_filter, wl.hkernel, wl.wkernel)\n",
    "    b_shape = (wl.batch, wl.out_filter, 1, 1)\n",
    "    if data_pack:\n",
    "        data_shape = (\n",
    "            wl.batch // env.BATCH,\n",
    "            wl.in_filter // env.BLOCK_IN,\n",
    "            wl.height,\n",
    "            wl.width,\n",
    "            env.BATCH,\n",
    "            env.BLOCK_IN,\n",
    "        )\n",
    "        kernel_shape = (\n",
    "            wl.out_filter // env.BLOCK_OUT,\n",
    "            wl.in_filter // env.BLOCK_IN,\n",
    "            wl.hkernel,\n",
    "            wl.wkernel,\n",
    "            env.BLOCK_OUT,\n",
    "            env.BLOCK_IN,\n",
    "        )\n",
    "        bias_shape = (\n",
    "            wl.batch // env.BATCH,\n",
    "            wl.out_filter // env.BLOCK_OUT,\n",
    "            1,\n",
    "            1,\n",
    "            env.BATCH,\n",
    "            env.BLOCK_OUT,\n",
    "        )\n",
    "    else:\n",
    "        data_shape = a_shape\n",
    "        kernel_shape = w_shape\n",
    "        bias_shape = b_shape\n",
    "    data = te.placeholder(data_shape, name=\"data\", dtype=env.inp_dtype)\n",
    "    kernel = te.placeholder(kernel_shape, name=\"kernel\", dtype=env.wgt_dtype)\n",
    "    bias = te.placeholder(bias_shape, name=\"bias\", dtype=env.acc_dtype)\n",
    "    padding = relay.nn.get_pad_tuple2d((wl.hpad, wl.wpad))\n",
    "\n",
    "    # Define base computation schedule\n",
    "    with target:\n",
    "        if data_pack:\n",
    "            res = conv2d_fcompute(\n",
    "                data, kernel, (wl.hstride, wl.wstride), padding, (1, 1), layout, env.acc_dtype\n",
    "            )\n",
    "        else:\n",
    "            res = conv2d_fcompute(\n",
    "                data, kernel, (wl.hstride, wl.wstride), padding, (1, 1), env.acc_dtype\n",
    "            )\n",
    "        res = topi.right_shift(res, 8)\n",
    "        res = topi.add(res, bias)\n",
    "        res = my_clip(res, 0, (1 << env.OUT_WIDTH - 1) - 1)\n",
    "        res = topi.cast(res, env.out_dtype)\n",
    "        # Derive base schedule\n",
    "        s = conv2d_fschedule([res])\n",
    "        if print_ir:\n",
    "            print(vta.lower(s, [data, kernel, bias, res], simple_mode=True))\n",
    "\n",
    "    # Derive number of ops\n",
    "    fout_height = (wl.height + 2 * wl.hpad - wl.hkernel) // wl.hstride + 1\n",
    "    fout_width = (wl.width + 2 * wl.wpad - wl.wkernel) // wl.wstride + 1\n",
    "    num_ops = (\n",
    "        2\n",
    "        * wl.batch\n",
    "        * fout_height\n",
    "        * fout_width\n",
    "        * wl.hkernel\n",
    "        * wl.wkernel\n",
    "        * wl.out_filter\n",
    "        * wl.in_filter\n",
    "    )\n",
    "\n",
    "    # @memoize(\"vta.tests.test_benchmark_topi.conv2d.verify_nhwc\")\n",
    "    def get_ref_data():\n",
    "        # derive min max for act, wgt, and bias types (max non inclusive)\n",
    "        a_min, a_max = 0 - (1 << (env.INP_WIDTH - 1)), (1 << (env.INP_WIDTH - 1))\n",
    "        w_min, w_max = 0 - (1 << (env.WGT_WIDTH - 1)), (1 << (env.WGT_WIDTH - 1))\n",
    "        b_min, b_max = 0 - 1 << (env.INP_WIDTH + env.WGT_WIDTH - 2), 1 << (\n",
    "            env.INP_WIDTH + env.WGT_WIDTH - 2\n",
    "        )\n",
    "        a_np = np.random.randint(a_min, a_max, size=a_shape).astype(data.dtype)\n",
    "        w_np = np.random.randint(w_min, w_max, size=w_shape).astype(kernel.dtype)\n",
    "        b_np = np.random.randint(b_min, b_max, size=b_shape).astype(env.acc_dtype)\n",
    "        r_np = tvm.topi.testing.conv2d_nchw_python(\n",
    "            a_np.astype(env.acc_dtype),\n",
    "            w_np.astype(env.acc_dtype),\n",
    "            (wl.hstride, wl.wstride),\n",
    "            wl.hpad,\n",
    "        ).astype(env.acc_dtype)\n",
    "        return a_np, w_np, b_np, r_np\n",
    "\n",
    "    # Data in original format\n",
    "    data_np, kernel_np, bias_np, res_ref = get_ref_data()\n",
    "    if data_pack:\n",
    "        data_np = data_np.reshape(\n",
    "            wl.batch // env.BATCH,\n",
    "            env.BATCH,\n",
    "            wl.in_filter // env.BLOCK_IN,\n",
    "            env.BLOCK_IN,\n",
    "            wl.height,\n",
    "            wl.width,\n",
    "        ).transpose((0, 2, 4, 5, 1, 3))\n",
    "        kernel_np = kernel_np.reshape(\n",
    "            wl.out_filter // env.BLOCK_OUT,\n",
    "            env.BLOCK_OUT,\n",
    "            wl.in_filter // env.BLOCK_IN,\n",
    "            env.BLOCK_IN,\n",
    "            wl.hkernel,\n",
    "            wl.wkernel,\n",
    "        ).transpose((0, 2, 4, 5, 1, 3))\n",
    "        bias_np = bias_np.reshape(\n",
    "            wl.batch // env.BATCH, wl.out_filter // env.BLOCK_OUT, 1, 1, env.BATCH, env.BLOCK_OUT\n",
    "        )\n",
    "\n",
    "    # Build\n",
    "    if \"vta\" in target.keys:\n",
    "        with vta.build_config(disabled_pass={\"tir.CommonSubexprElimTIR\"}):\n",
    "            mod = vta.build(\n",
    "                s,\n",
    "                [data, kernel, bias, res],\n",
    "                target=tvm.target.Target(target, host=env.target_host),\n",
    "                name=\"conv2d\",\n",
    "            )\n",
    "    else:\n",
    "        mod = tvm.build(\n",
    "            s,\n",
    "            [data, kernel, bias, res],\n",
    "            target=tvm.target.Target(target, host=env.target_host),\n",
    "            name=\"conv2d\",\n",
    "        )\n",
    "    temp = utils.tempdir()\n",
    "    mod.save(temp.relpath(\"conv2d.o\"))\n",
    "    remote.upload(temp.relpath(\"conv2d.o\"))\n",
    "    f = remote.load_module(\"conv2d.o\")\n",
    "    dev = remote.device(str(target))\n",
    "\n",
    "    res_np = np.zeros(topi.utils.get_const_tuple(res.shape)).astype(res.dtype)\n",
    "    data_arr = tvm.nd.array(data_np, dev)\n",
    "    kernel_arr = tvm.nd.array(kernel_np, dev)\n",
    "    bias_arr = tvm.nd.array(bias_np, dev)\n",
    "    res_arr = tvm.nd.array(res_np, dev)\n",
    "    time_f = f.time_evaluator(\"conv2d\", dev, number=samples)\n",
    "\n",
    "    # In vta sim mode, collect simulator runtime statistics\n",
    "    stats = {}\n",
    "    cost = None\n",
    "    if env.TARGET in [\"sim\", \"tsim\"]:\n",
    "        # Check if we're in local RPC mode (allows us to rebuild the\n",
    "        # runtime on the fly when varying the VTA designs)\n",
    "        local_rpc = int(os.environ.get(\"VTA_LOCAL_SIM_RPC\", \"0\"))\n",
    "        if local_rpc:\n",
    "            if env.TARGET == \"sim\":\n",
    "                remote.get_function(\"vta.simulator.profiler_clear\")()\n",
    "            else:\n",
    "                remote.get_function(\"vta.tsim.profiler_clear\")()\n",
    "            cost = time_f(data_arr, kernel_arr, bias_arr, res_arr)\n",
    "            if env.TARGET == \"sim\":\n",
    "                stats = json.loads(remote.get_function(\"vta.simulator.profiler_status\")())\n",
    "            else:\n",
    "                stats = json.loads(remote.get_function(\"vta.tsim.profiler_status\")())\n",
    "        else:\n",
    "            simulator.clear_stats()\n",
    "            cost = time_f(data_arr, kernel_arr, bias_arr, res_arr)\n",
    "            stats = simulator.stats()\n",
    "    else:\n",
    "        cost = time_f(data_arr, kernel_arr, bias_arr, res_arr)\n",
    "\n",
    "    # Check correctness\n",
    "    correct = False\n",
    "    if check_correctness:\n",
    "        res_orig = res_arr.numpy()\n",
    "        if data_pack:\n",
    "            res_orig = res_orig.transpose((0, 4, 1, 5, 2, 3)).reshape(\n",
    "                wl.batch, wl.out_filter, fout_height, fout_width\n",
    "            )\n",
    "            bias_np = bias_np.transpose((0, 4, 1, 5, 2, 3)).reshape(wl.batch, wl.out_filter, 1, 1)\n",
    "        res_ref = res_ref >> env.WGT_WIDTH\n",
    "        res_ref += bias_np\n",
    "        res_ref = np.clip(res_ref, 0, (1 << env.OUT_WIDTH - 1) - 1)\n",
    "        res_ref = res_ref.astype(env.out_dtype)\n",
    "        correct = np.allclose(res_orig, res_ref)\n",
    "\n",
    "    gops = (num_ops / cost.mean) / float(10**9)\n",
    "    status = \"PASSED\" if correct else \"FAILED\"\n",
    "    if \"arm_cpu\" in target.keys:\n",
    "        device = \"CPU\"\n",
    "    elif \"vta\" in target.keys:\n",
    "        device = \"VTA\"\n",
    "    print(\"%s CONV2D TEST %s: Time cost = %g sec/op, %g GOPS\" % (device, status, cost.mean, gops))\n",
    "\n",
    "    return correct, cost, stats\n",
    "\n",
    "\n",
    "@pytest.mark.parametrize(\"device\", [\"vta\", \"arm_cpu\"])\n",
    "def test_conv2d(device):\n",
    "    def _run(env, remote):\n",
    "        if device == \"vta\":\n",
    "            target = env.target\n",
    "            if env.TARGET not in [\"sim\", \"tsim\", \"intelfocl\"]:\n",
    "                assert tvm.runtime.enabled(\"rpc\")\n",
    "                program_fpga(remote, bitstream=None)\n",
    "                reconfig_runtime(remote)\n",
    "        elif device == \"arm_cpu\":\n",
    "            target = env.target_vta_cpu\n",
    "        with autotvm.tophub.context(target):  # load pre-tuned schedule parameters\n",
    "            for _, wl in resnet_wkls:\n",
    "                print(wl)\n",
    "                run_conv2d(env, remote, wl, target)\n",
    "\n",
    "    vta.testing.run(_run)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    test_conv2d(device=\"arm_cpu\")\n",
    "    test_conv2d(device=\"vta\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
